"""
Original GCR-NCCL Runner Implementation.

This runner executes the original GCR solvers using mpirun.
"""

import os
import subprocess
import time
import logging
from pathlib import Path
from typing import Dict, List, Optional, Any

from gcr_solver_manager.extensions.base_runner import (
    RunnerInterface, RunConfig, RunResult, EnvironmentResult
)

logger = logging.getLogger(__name__)


class OriginalGCRRunner(RunnerInterface):
    """Runner for Original GCR-NCCL framework."""
    
    def __init__(self, source_dir: Path):
        """Initialize runner with source directory."""
        self.source_dir = source_dir
    
    def run_solver(self, run_config: RunConfig) -> RunResult:
        """Execute solver with given configuration."""
        start_time = time.time()
        
        try:
            # Build mpirun command
            cmd = self._build_mpirun_command(run_config)
            
            logger.info(f"Executing solver: {' '.join(cmd)}")
            
            # Set up environment
            env = os.environ.copy()
            env.update(run_config.environment)
            
            # Execute solver
            result = subprocess.run(
                cmd,
                cwd=self.source_dir,
                capture_output=True,
                text=True,
                env=env,
                timeout=run_config.timeout or 1800  # Default 30 minutes
            )
            
            duration = time.time() - start_time
            success = result.returncode == 0
            
            # Extract basic metrics from output
            metrics = self._extract_metrics(result.stdout, result.stderr)
            
            run_result = RunResult(
                success=success,
                exit_code=result.returncode,
                stdout=result.stdout,
                stderr=result.stderr,
                duration=duration,
                log_file=None,  # No specific log file for original GCR
                metrics=metrics
            )
            
            logger.info(f"Solver execution completed in {duration:.1f}s, success={success}")
            return run_result
            
        except subprocess.TimeoutExpired:
            return RunResult(
                success=False,
                exit_code=-1,
                stdout="",
                stderr="Solver execution timed out",
                duration=time.time() - start_time,
                log_file=None,
                metrics={}
            )
        except Exception as e:
            return RunResult(
                success=False,
                exit_code=-1,
                stdout="",
                stderr=f"Solver execution failed: {e}",
                duration=time.time() - start_time,
                log_file=None,
                metrics={}
            )
    
    def check_environment(self) -> EnvironmentResult:
        """Validate runtime environment."""
        missing_components = []
        env_info = {}
        
        try:
            # Check for mpirun
            result = subprocess.run(['mpirun', '--version'], capture_output=True, text=True, timeout=10)
            if result.returncode == 0:
                env_info['mpirun'] = result.stdout.split('\n')[0]
            else:
                missing_components.append('mpirun')
            
            # Check for CUDA runtime (if GPU backend)
            result = subprocess.run(['nvidia-smi'], capture_output=True, text=True, timeout=10)
            if result.returncode == 0:
                env_info['nvidia_driver'] = 'available'
            else:
                env_info['nvidia_driver'] = 'not_available'
            
            # Check for HIP runtime (if HIP backend)
            result = subprocess.run(['rocm-smi'], capture_output=True, text=True, timeout=10)
            if result.returncode == 0:
                env_info['rocm'] = 'available'
            else:
                env_info['rocm'] = 'not_available'
            
            valid = len(missing_components) == 0
            
            return EnvironmentResult(
                valid=valid,
                missing_components=missing_components,
                environment_info=env_info,
                error_message=None if valid else f"Missing components: {', '.join(missing_components)}"
            )
            
        except Exception as e:
            return EnvironmentResult(
                valid=False,
                missing_components=['unknown'],
                environment_info={},
                error_message=f"Environment check failed: {e}"
            )
    
    def prepare_environment(self, run_config: RunConfig) -> bool:
        """Prepare environment for execution."""
        try:
            # Check if binary exists
            if not run_config.binary_path.exists():
                logger.error(f"Binary not found: {run_config.binary_path}")
                return False
            
            # Make sure binary is executable
            run_config.binary_path.chmod(0o755)
            
            return True
            
        except Exception as e:
            logger.error(f"Failed to prepare environment: {e}")
            return False
    
    def cleanup_environment(self, run_config: RunConfig) -> bool:
        """Cleanup after execution."""
        # No specific cleanup needed for original GCR
        return True
    
    def _build_mpirun_command(self, run_config: RunConfig) -> List[str]:
        """Build the mpirun command for executing the solver."""
        cmd = [
            "mpirun",
            "-np", str(run_config.mpi_processes),
            str(run_config.binary_path),
            f"--solver={run_config.solver_type}"
        ]
        
        # Add solver-specific flags
        if run_config.solver_type == 'ca-gcr':
            cmd.append("--use-ca")
        elif run_config.solver_type == 'gmres':
            cmd.append("--use-gmres")
        elif run_config.solver_type == 'bca-gmres':
            cmd.extend(["--use-gmres", "--use-ca"])
        
        # Add debug mode
        if run_config.debug_mode:
            cmd.append("--maxit=1")
        
        # Add extra arguments
        cmd.extend(run_config.extra_args)
        
        return cmd
    
    def _extract_metrics(self, stdout: str, stderr: str) -> Dict[str, Any]:
        """Extract performance and convergence metrics from solver output."""
        metrics = {}
        
        try:
            # Look for convergence information
            lines = stdout.split('\n') + stderr.split('\n')
            
            for line in lines:
                line = line.strip().lower()
                
                # Extract final residual norm
                if 'final' in line and 'residual' in line and 'norm' in line:
                    parts = line.split()
                    for i, part in enumerate(parts):
                        try:
                            if 'e-' in part or 'e+' in part or ('.' in part and part.replace('.', '').replace('-', '').isdigit()):
                                metrics['final_residual_norm'] = float(part)
                                break
                        except:
                            continue
                
                # Extract iteration count
                if 'iteration' in line and ':' in line:
                    parts = line.split(':')
                    if len(parts) > 1:
                        try:
                            metrics['iterations'] = int(parts[1].strip())
                        except:
                            pass
                
                # Extract convergence status
                if 'converged' in line:
                    metrics['converged'] = True
                elif 'not converged' in line or 'failed to converge' in line:
                    metrics['converged'] = False
            
        except Exception as e:
            logger.debug(f"Failed to extract metrics: {e}")
        
        return metrics